Enable offloading of validation models to CPU during training runs for EMTTS - #16035
Enable offloading of validation models to CPU during training runs for EMTTS#16035blisc wants to merge 1 commit into
Conversation
Signed-off-by: Jason <jasoli@nvidia.com>
| ) | ||
|
|
||
| def training_step(self, batch, batch_idx): | ||
| uses_uncached_audio = 'context_audio_codes' not in batch or 'audio_codes' not in batch |
There was a problem hiding this comment.
We should extend this if statement to check whether the codec is required for multi-turn user-audio conditioning:
self.cfg.get("use_multiturn_dataset", False)
and batch["user_audio_turn_splitted"] is not None
and self.cfg.get("condition_on_user_speech", False)
Without this check, when context_audio_codes and audio_codes are included in the multi-turn data, user-audio conditioning feature extraction may run with the codec on the CPU, which would be extremely slow or raise errors.
| audio_lens = batch['audio_lens'] | ||
| audio_codes, audio_codes_lens = self._codec_helper.audio_to_codes(audio, audio_lens) | ||
|
|
||
| if uses_uncached_audio and self._should_offload_validation_models(): |
There was a problem hiding this comment.
Why is this necessary? Could we keep the codec on the GPU throughout an epoch whenever any batch in that epoch requires it?
If you still want to do it, please move it to the after the following block, otherwise multiturn training data will fails.
if (
self.cfg.get("use_multiturn_dataset", False)
and batch["user_audio_turn_splitted"] is not None
and self.cfg.get("condition_on_user_speech", False)
):
Edresson
left a comment
There was a problem hiding this comment.
Why is this change necessary?
Also, have you measured its impact on validation performance? Validation is already relatively slow, and I’m concerned that this PR may introduce additional overhead.
What does this PR do ?
Enable offloading of validation models to CPU during training runs for EMTTS
Collection: tts
Changelog
PR Type: